# -*- coding: utf-8 -*-
# @Time    : 2023/2/21 17:29
# @Author  : HeAlec
# @FileName: rlts.py
# @Desc: description
# @blogs ：https://segmentfault.com/u/alec_5e8962e4635ca
import time
import numpy as np
import os
from RLOnline.rl_brain import PolicyGradient
from RLOnline.rl_env_inc import TrajComp


class RLAgent:
    def __init__(self, datasets, dataSize, metric):
        # todo 轨迹文件的路径
        self.traj_path = datasets
        if metric == "ss":
            metric = "ped"

        a_size = 3  # RLTS 3, RLTS-Skip 5
        s_size = 3  # RLTS and RLTS-Skip are both 3 online
        self.env = TrajComp(self.traj_path, dataSize, a_size, s_size, metric)
        self.RL = PolicyGradient(self.env.n_features, self.env.n_actions)
        self.RL.load('./saved_rlts/save/0.00051169259094067_ratio_0.1/')

    def RL_online(self, buffer_size, episode):
        if buffer_size < 3:
            return
        steps, observation = self.env.reset(episode, buffer_size)
        tic1 = time.perf_counter()
        for index in range(buffer_size, steps):
            if index == steps - 1:
                done = True
            else:
                done = False
            action = self.RL.quick_time_action(observation)  # matrix implementation for fast efficiency when the model is ready
            observation_, _ = self.env.step(episode, action, index, done,
                                            'V')  # 'T' means Training, and 'V' means Validation
            observation = observation_
        tic2 = time.perf_counter()
        idx, max_err = self.env.output(episode, 'V')
        if idx[-1] == steps:
            idx, max_err = self.env.output(episode, 'V')
        tm = (tic2 - tic1) / len(self.env.ori_traj_set[episode])
        return None, idx, max_err, tm

    def run_all(self, size, ratio, metric):
        err = []
        timelist = []
        for i in range(size):
            buffer_size = int(ratio * len(self.env.ori_traj_set[i]))
            _, idx, max_err, tm = self.RL_online(buffer_size, i)
            timelist.append(tm)
            err.append(max_err)
        print(f"rlts压缩率 {ratio}, 耗时 {np.mean(timelist)}, error {np.mean(err)}")